import Router from '@koa/router'
import { DocxLoader } from '@langchain/community/document_loaders/fs/docx'
import { PDFLoader } from '@langchain/community/document_loaders/fs/pdf'
import { UnstructuredLoader } from '@langchain/community/document_loaders/fs/unstructured'
import {
    Chroma,
    ChromaDeleteParams,
} from '@langchain/community/vectorstores/chroma'
import {
    MarkdownTextSplitter,
    RecursiveCharacterTextSplitter,
} from '@langchain/textsplitters'
import { TextLoader } from 'langchain/document_loaders/fs/text'
import fs from 'node:fs'
import path from 'node:path'
import KnowledgeBaseFile from '../../models/KnowledgeBaseFile'
import { KnowledgeBaseFileStatus } from './../../models/KnowledgeBaseFile'
import { getVectorStore } from './vectorStore'
import sequelize from '../../models/db'
import { toFilepath } from '../../utils/uploadUrlTrans'

const typeLoaderMap = {
    txt: TextLoader,
    md: UnstructuredLoader,
    pdf: PDFLoader,
    docx: DocxLoader,
}

const textSplitterMap: {
    [key in keyof typeof typeLoaderMap]:
        | typeof RecursiveCharacterTextSplitter
        | typeof MarkdownTextSplitter
} = {
    txt: RecursiveCharacterTextSplitter,
    md: MarkdownTextSplitter,
    pdf: RecursiveCharacterTextSplitter,
    docx: RecursiveCharacterTextSplitter,
}

// 把文件向量化，保存到向量数据库
export const saveToVectorStore = async (data: {
    knowledgeBaseFileId: number
    knowledgeBaseId: number
    collectionName: string
    url: string
    type: keyof typeof typeLoaderMap
}) => {
    const { type, knowledgeBaseFileId, knowledgeBaseId, collectionName, url } =
        data
    try {
        // 向量存储，向量化模型，collection
        const vectorStore = await getVectorStore(collectionName)
        // 文件路径
        const filePath = path.resolve('./public', toFilepath(url))
        // 加载文件
        const Loader = typeLoaderMap[type]
        if (!Loader) {
            throw new Error('未找到loader')
        }
        // 加载文档
        const loader = new Loader(filePath)
        const documents = await loader.load()
        // 拆分文件
        const DocSplitter = textSplitterMap[type]
        const docSplitter = new DocSplitter()
        const texts = await docSplitter.splitDocuments(documents)
        // 处理metadata
        texts.forEach((doc) => {
            doc.metadata = {
                ...doc.metadata,
                type,
                knowledgeBaseFileId,
                knowledgeBaseId,
                collectionName,
                url,
            }
        })
        // 更新数据库为正在向量化
        KnowledgeBaseFile.update(
            {
                status: KnowledgeBaseFileStatus.Embedding,
                metaData: sequelize.literal(
                    `JSON_SET(metaData, '$.docCount', ${texts.length})`
                ),
            },
            {
                where: { id: knowledgeBaseFileId },
            }
        )
        // 保存到向量化存储的collection
        const stored = await vectorStore.addDocuments(texts)
        // 更新数据库为向量化完成
        KnowledgeBaseFile.update(
            {
                status: KnowledgeBaseFileStatus.Embedded,
                metaData: sequelize.literal(
                    `JSON_SET(metaData, '$.docIds', '${JSON.stringify(
                        stored
                    )}')`
                ),
            },
            {
                where: { id: knowledgeBaseFileId },
            }
        )
    } catch (error) {
        console.log(error)
        KnowledgeBaseFile.update(
            {
                status: KnowledgeBaseFileStatus.Failed,
            },
            {
                where: { id: knowledgeBaseFileId },
            }
        ).catch(() => {})
    }
}

// 把文件从向量化数据库中删除
export const delFromVectorStore = async (
    collectionName: string,
    delFuncParams: ChromaDeleteParams<Chroma['FilterType']>
) => {
    // 向量存储，向量化模型，collection
    const vectorStore = await getVectorStore(collectionName)
    // 删除
    await vectorStore.delete(delFuncParams)
}

export const delKnowledgeBaseFile: Router.Middleware = async (ctx) => {
    const { id, knowledgeBaseId, collectionName, url } = ctx.request.body

    if (!id || !knowledgeBaseId || !collectionName || !url) {
        ctx.body = { success: false, msg: '参数错误' }
        return
    }

    try {
        // 从向量化数据库中删除
        await delFromVectorStore(collectionName, {
            filter: {
                knowledgeBaseFileId: id,
            },
        })

        // 从本地删除
        const filePath = path.resolve('./public', toFilepath(url))
        fs.unlinkSync(filePath)

        // 从数据库中删除
        const list = await KnowledgeBaseFile.destroy({
            where: { id },
        })

        ctx.body = { success: true, data: list }
    } catch (error) {
        console.log(error)
        ctx.body = { success: false, msg: '删除失败' }
    }
}
